package org.ws.jboss.security.spi;

import java.net.UnknownHostException;
import java.security.Principal;
import java.security.acl.Group;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

import javax.security.auth.Subject;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.login.FailedLoginException;
import javax.security.auth.login.LoginException;

import org.jboss.security.SimpleGroup;
import org.jboss.security.Util;
import org.jboss.security.auth.spi.UsernamePasswordLoginModule;

import com.mongodb.BasicDBList;
import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.MongoClient;
import com.mongodb.MongoClientURI;

public class MongoDBLoginModule extends UsernamePasswordLoginModule{
	
	   protected String mongoDBConnectionURI;
	   protected String dbName;
	   protected String principalsCollection = "users";
	   protected String userNameField = "username";
	   protected String passwordField = "password";
	   protected String rolesCollection = "roles";
	   protected String rolesFieldName = "roles";
	   protected String rolesGroupeFieldName = "roleGroup";
	   protected String rolesUsernameField = "username";
	   protected String roleNameField = "role";
	   
	   protected String dbUserName;
	   protected String dbPassword;
	   
	   /**
	    * Initialize this LoginModule.
	    * 
	    * @param options -
	    * mongoDBConnectionURI: The connection URI for the MongDB server
	    * principalsCollection: The collection that holds the principals
	    * principalsQuery: MongoDB query to locate the user
	    * rolesCollection: The collection that holds the roles
	    * rolesQuery: The MongoDB query to locate the users roles
	    */
	   @SuppressWarnings({ "unchecked", "rawtypes" })
	public void initialize(Subject subject, CallbackHandler callbackHandler,
	      Map sharedState, Map options)
	   {
	      super.initialize(subject, callbackHandler, sharedState, options);
	      
	      // Connection URI of the db including the port number e.g. mongodb://localhost:27017
	      mongoDBConnectionURI = (String) options.get("mongDBConnectionURI");
	      if( mongoDBConnectionURI == null )
	    	  mongoDBConnectionURI = "mongodb://localhost:27017/";
	      
	      //nome of the MongoDB database
	      Object tmp = options.get("databaseName");
	      if(tmp != null) dbName = tmp.toString();
	      
	      //name of the collection that holds the principals
	      tmp = options.get("principalsCollection");
	      if(tmp != null) principalsCollection = tmp.toString();
	      
	      //name of the filed in the principalsCollection that holds the username
	      tmp = options.get("userNameField");
	      if( tmp != null )
	         userNameField = tmp.toString();
	      
	      //name of the field in the principalsCollection that holds the users password
	      tmp = options.get("passwordField");
	      if(tmp != null) passwordField = tmp.toString();
	      
	      //name of the collection that holds the roles.  Can be the same as the users collection.
	      tmp = options.get("rolesCollection");
	      if(tmp != null) rolesCollection = tmp.toString();
	      
	    //field name in the roles collection that holds the username.
	      tmp = options.get("rolesUsernameField");
	      if(tmp != null) rolesUsernameField = tmp.toString();
	      
	      /*name of the field that holds the users roles
	       * this must return an array.
	       * 
	       * e.g. 	{"username":"xxxx",
	       * 		"roles":[{"role":"users", "group":"Roles"},{"role":"admins", "group":"Roles"}]
	       * 		}
	       */
	      tmp = options.get("rolesFieldName");
	      if(tmp != null) rolesFieldName = tmp.toString();
	      
	      //field name in the roles collection that holds the gorup name.  This should always be Roles with a capital R
	      tmp = options.get("rolesGroupFieldName");
	      if(tmp != null) rolesGroupeFieldName = tmp.toString();
	          
	      //field name in the roles collection that holds the name of the role
	      tmp = options.get("roleNameField");
	      if(tmp != null) roleNameField = tmp.toString();
	      
	      //username used to log into the database
	      tmp = options.get("databaseUser");
	      if(tmp != null) dbUserName = tmp.toString();
	      
	      //password used to log into the datbase
	      tmp = options.get("databasePassword");
	      if(tmp != null) dbPassword = tmp.toString();
	      
	      
	      log.trace("DatabaseServerLoginModule, mongoDBCollection = "+mongoDBConnectionURI);
	      log.trace("principalsCollection = " + principalsCollection);
	      log.trace("userNameField ="+userNameField);
	      log.trace("passwordField = " + passwordField);
	      log.trace("rolesCollection=" + rolesCollection);
	      log.trace("rolesUsernameField = " + rolesUsernameField);
	      log.trace("rolesFieldName = " + rolesFieldName);
	      log.trace("rolesGroupFieldName = " + rolesGroupeFieldName);
	      log.trace("roleNameField = " + roleNameField);
	      log.trace("databaseUser = " + dbUserName);
	      if(dbPassword != null){
	    	  log.trace("databasePassword = '*************'");
	      }else{
	    	  log.trace("databasePassword = " + dbPassword);
	      }	      
	   }

	   /** Get the expected password for the current username available via
	    * the getUsername() method. This is called from within the login()
	    * method after the CallbackHandler has returned the username and
	    * candidate password.
	    * @return the valid password String
	    */
	   protected String getUsersPassword() throws LoginException
	   {
	      String username = getUsername();
	      String password = null;
	      if(dbName == null || dbName.length() == 0) throw new LoginException("No databaseName configured");
	      MongoClient mongoClient = null;
	      
	      try
	      {
	    	  //connect to the mongoDatabase
	    	  mongoClient = new MongoClient(new MongoClientURI(mongoDBConnectionURI));
	    	  	    	  
	    	  DB db = mongoClient.getDB(dbName);
		      
		      boolean auth = db.authenticate(dbUserName, dbPassword.toCharArray());
		      
		      if(!auth) throw new LoginException("Failed to connect to the database (authentication)");
		      
	         DBCollection usersColl = db.getCollection(principalsCollection);
	         
	         if(usersColl == null) throw new LoginException("Could not find principals collection - " + principalsCollection);
	         BasicDBObject userRecord = (BasicDBObject) usersColl.findOne(new BasicDBObject(userNameField, username));
	         
	         if(userRecord == null) throw new FailedLoginException("No matching username found in Principals");
	        	           
	         password = userRecord.getString(passwordField);
	         if(password == null || password.length() == 0) throw new LoginException("No password field found");
	         password = convertRawPassword(password);
	      } catch (UnknownHostException e) {
			throw new LoginException("Could not connect to mongo server @ " + mongoDBConnectionURI);
		}
	      finally
	      {
	    	  if(mongoClient != null)  mongoClient.close();
	      }
	      return password;
	   }

	   /** Execute the rolesQuery against the dsJndiName to obtain the roles for
	    the authenticated user.
	     
	    @return Group[] containing the sets of roles
	    */
	   protected Group[] getRoleSets() throws LoginException
	   {
	      String username = getUsername();
	      Group[] roleSets;
	      HashMap<String, Group> setsMap = new HashMap<String, Group>();
	      if(dbName == null || dbName.length() == 0) throw new LoginException("No databaseName configured");
	      MongoClient mongoClient = null;
	      try{
		      mongoClient = new MongoClient(new MongoClientURI(mongoDBConnectionURI));
		      DB db = mongoClient.getDB(dbName);
		      boolean auth = db.authenticate(dbUserName, dbPassword.toCharArray());
		      
		      if(!auth) throw new LoginException("Failed to connect to the database (authentication)");
		      
		      DBCollection rolesColl = db.getCollection(rolesCollection);
		      
		      if(rolesColl == null) throw new LoginException("Cound not find rolesCollection - " + rolesCollection);
	         
		      BasicDBObject userRecord = (BasicDBObject) rolesColl.findOne(new BasicDBObject(rolesUsernameField, username));
		      
		      if(userRecord == null) throw new LoginException("No user found in roles collection");
		      
		      @SuppressWarnings("unchecked")
			ArrayList<BasicDBObject> roles = (ArrayList<BasicDBObject>)userRecord.get(rolesFieldName);
		      
		      if(roles == null || roles.size() == 0){
		    	  if(this.getUnauthenticatedIdentity() == null){
		    		  throw new FailedLoginException("No matching username found in Roles");
		    	  }
		    	  roleSets = new Group[1];
		    	  roleSets[0] = new SimpleGroup("Roles");
		    	  return roleSets;
		      }
		      
		      for(BasicDBObject role : roles){
		    	  String groupName = role.getString(rolesGroupeFieldName);
		    	  String roleName = role.getString(roleNameField);
		    	  
		    	  if(groupName == null || groupName.length() == 0){
		    		  groupName = "Roles";
		    	  }
		    	  
		    	  Group group = (Group)setsMap.get(groupName);
		    	  if(group == null){
		    		  group = new SimpleGroup(groupName);
		    		  setsMap.put(groupName, group);
		    	  }
		    	  
		    	  try{
		    		  Principal p = this.createIdentity(roleName);
		    		  log.trace("Assign user to role " + roleName);
		    		  System.out.println("Assign user to role " + roleName);
		    		  group.addMember(p);
		    	  }catch(Exception e)
		            {
		               log.debug("Failed to create principal: "+roleName, e);
		               System.out.println("Failed to create principal: "+roleName);
		               e.printStackTrace();
		            }
		    		  
		      }
		      
	      } catch (UnknownHostException e) {
				throw new LoginException("Could not connect to mongo server @ " + mongoDBConnectionURI);
			}
		      finally
		      {
		    	  if(mongoClient != null)  mongoClient.close();
		      }
	      roleSets = new Group[setsMap.size()];
	      setsMap.values().toArray(roleSets);
	      return roleSets;
	   }
	   
	   /** A hook to allow subclasses to convert a password from the database
	    into a plain text string or whatever form is used for matching against
	    the user input. It is called from within the getUsersPassword() method.
	    @param rawPassword - the password as obtained from the database
	    @return the argument rawPassword
	    */
	   protected String convertRawPassword(String rawPassword)
	   {
	      return rawPassword;
	   }
}
